package zy.learn.demo.structuredstreaming.window

import java.sql.Timestamp

import org.apache.spark.SparkConf
import org.apache.spark.sql.SparkSession
import org.apache.spark.sql.streaming.OutputMode

object WordCountWindow {
  def main(args: Array[String]): Unit = {
    val sparkConf = new SparkConf().set("spark.sql.shuffle.partitions", "3")
    val spark = SparkSession.builder()
      .master("local[2]")
      .config(sparkConf)
      .appName("WordCountWindow")
      .getOrCreate()

    import spark.implicits._

    val lines = spark.readStream
      .format("socket") // 设置数据源
      .option("host", "co7-203")
      .option("port", 9999)
      .option("includeTimestamp", true) // 给产生的数据自动添加时间戳，这个时间戳是固定值
      .load
    // 把行切割成单词, 保留时间戳，时间戳类型为 java.sql.Timestamp ！！！
    /*val wordsDF = lines.as[(String, TimestampType)].flatMap(line => {
      line._1.split(" ").map((_, line._2))
    }).toDF("word", "timestamp")*/
    val wordsDF = lines.as[(String, Timestamp)].flatMap({
      case (words, ts) => words.split("\\W+").map((_, ts))
    }).toDF("word", "ts")

    import org.apache.spark.sql.functions._

    val wordCounts = wordsDF.groupBy(
      // 调用 window 函数, 返回的是一个 Column 参数 1: df 中表示时间戳的列 参数 2: 窗口长度 参数 3: 滑动步长
      window($"ts", "10 seconds", "5 seconds"),
      $"word"
    ).count().orderBy($"window") // 计数，并按照窗口排序

    val query = wordCounts.writeStream
      .outputMode(OutputMode.Complete())
      .format("console")
      .option("truncate", "false")
      .start()

    query.awaitTermination()
  }
}
